# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: chunhui

import collections
import datetime
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd

from absl.flags import FLAGS
from python.coin.pnl.sim_stat_frompb import axvlines


def join_df_by_time(sim_df, live_df, frequency, keep_columns=None):
  joined_df = pd.DataFrame()
  for pnldf in (sim_df, live_df):
    pnldf.index = pd.to_datetime(pnldf.timestamp).dt.floor(freq=frequency)
    df = pnldf.loc[~pnldf.index.duplicated(keep='last')].copy()
    if keep_columns:
      df = df[keep_columns]
    if joined_df.empty:
      joined_df = df
    else:
      joined_df = pd.merge(joined_df, df, left_index=True, right_index=True, how='outer')
    pnldf = pnldf.reset_index(drop=True)
  return joined_df


def compute_correlations(df, columns):
  corr = {}
  for column in columns:
    for id in ('x', 'y'):
      df[f'{column}_diff_{id}'] = df[f'{column}_{id}'].fillna(method='ffill') \
                                  .fillna(method='bfill').diff().fillna(0)

    corr[column] = df[f'{column}_diff_x'].corr(df[f'{column}_diff_y'])
  return corr


def format_stats(x):
  if isinstance(x, float) and not np.isnan(x) and abs(x) > 0:
    accuracy = np.clip(np.ceil(-np.log10(abs(x)) + 4), 2, 6)
    if not np.isnan(accuracy):
      return "{:,}".format(round(x, int(accuracy)))
  return x


def dump_stats(df, correlation, sim_key, live_key, output_dir, filename):
  sim_symbol = ".".join(sim_key.split('.')[1:]).replace('CURRENT_', '')
  live_symbol = live_key.replace('CURRENT_', '')
  extra = collections.defaultdict(list)
  for key, corr in correlation.items():
    extra['Type'].append(f'{key}_correlation')
    extra['Sim'].append(corr)
    extra['Live'].append(corr)
  extra['Type'].append('symbol')
  extra['Sim'].append(sim_symbol)
  extra['Live'].append(live_symbol)

  extra_df = pd.DataFrame(extra).set_index('Type').T.reset_index(drop=True)
  df = pd.concat([df, extra_df], sort=False, axis=1)
  df = df.rename(columns=lambda x: x.split('(')[0])
  df.to_csv(os.path.join(output_dir, f'{filename}.csv'), index=False)


def plot_graph(sim_df, live_df, name, sim_stat, live_stat, sim_dti, live_dti,
               sim_key, live_key, output_dir, xlim=None):
  os.makedirs(output_dir, exist_ok=True)
  plt.rcParams["figure.figsize"] = 14, 12
  keys = [f'sim {sim_key}', f'live {live_key}']
  colors = ['tab:blue', 'tab:orange']
  dti = live_dti
  dt = dti.date
  datesplit = np.where(np.diff(dt) > datetime.timedelta(0))[0]

  plt.subplot(211)  # net pnl

  for i, pnldf in enumerate((sim_df, live_df)):
    plt.plot(pnldf['tindex'],
             pnldf['pnl_net'],
             '-',
             color=colors[i],
             markersize=1,
             linewidth=0.5,
             drawstyle='steps-post')
  if xlim is not None:
    plt.xlim(xlim[0], xlim[1])
  plt.ylabel('Net PNL')
  plt.grid(True)

  axvlines(dti[datesplit], color='b', lw=0.2)
  plt.legend(keys, loc='upper left')

  plt.subplot(337)  # volume
  for i, pnldf in enumerate((sim_df, live_df)):
    plt.plot(pnldf['tindex'],
             pnldf['fill_cash_cumsum'],
             '-',
             color=colors[i],
             markersize=1,
             linewidth=0.5,
             drawstyle='steps-post')
  if xlim is not None:
    plt.xlim(xlim[0], xlim[1])
  axvlines(dti[datesplit], color='b', lw=0.2)
  plt.xticks(rotation=40)
  plt.legend(keys, loc='upper left')
  plt.ylabel('Volume')
  plt.grid(True)

  live_sharpe = live_df.pnl_net.diff().mean() / live_df.pnl_net.diff().std()
  sim_sharpe = sim_df.pnl_net.diff().mean() / sim_df.pnl_net.diff().std()

  correlation_columns = ['pnl_net', 'maker_fill_cumsum', 'taker_fill_cumsum']
  hourly_df = join_df_by_time(sim_df, live_df, '3600s', correlation_columns)
  correlation = compute_correlations(hourly_df, correlation_columns)

  plt.subplot(3, 3, (8, 9))  # sharpe, correlation

  stats = pd.DataFrame({'Sim': sim_stat, 'Live': live_stat}).T
  stats = stats.reset_index().rename(columns={'index': 'Type'})
  stats['sharpe'] = pd.Series([sim_sharpe, live_sharpe])
  stats['Sim-Live correlation'] = pd.Series([correlation['pnl_net'], np.nan])

  if FLAGS.dump_stats:
    dump_stats(stats.copy(), correlation, sim_key, live_key, output_dir, name)
  stats = stats.fillna('').set_index('Type').T

  for key in ['Sim', 'Live']:
    stats[key] = stats[key].map(lambda x: format_stats(x))
  tbl = plt.table(stats.values,
                  colWidths=[0.2, 0.2],
                  fontsize=14,
                  rowLabels=stats.index,
                  colLabels=stats.columns,
                  loc="center")
  tbl.scale(1.6, 1.3)
  plt.axis('off')

  plt.subplot(614)  # Position
  for i, pnldf in enumerate((sim_df, live_df)):
    plt.plot(pnldf['tindex'],
             pnldf['open_cash'],
             '-',
             color=colors[i],
             markersize=1,
             linewidth=0.5,
             drawstyle='steps-post')
  if xlim is not None:
    plt.xlim(xlim[0], xlim[1])
  axvlines(dti[datesplit], color='b', lw=0.2)
  plt.grid(True)
  plt.ylabel('Position')

  name = f"corr_{np.round(correlation['pnl_net'], 3):.3f}_" + name
  plt.suptitle(name)
  pic_path = os.path.join(output_dir, f'{name}_{FLAGS.trading_date}.png')
  plt.savefig(pic_path)
  plt.clf()

  sim_pnl, live_pnl = sim_df.iloc[-1]['pnl_net'], live_df.iloc[-1]['pnl_net']
  poor_performance = sim_pnl - live_pnl > min(abs(sim_pnl), abs(live_pnl))
  return pic_path, correlation['pnl_net'], poor_performance
